///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file matrix_math.cpp
///
/// \brief   implementation for matrix mathematics functions
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <src/pcps.h>
#include <src/matrix_math.h>
#include <src/blas_lapack.h>
#include <src/mantissa_exponent.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Computes a matrix's inverse and determinant using an LU decomposition.
///
/// \param[in]     n                       the dimension of the matrix A
/// \param[out]    det                     on exit, the matrix's determinant
/// \param[in,out] A          size n*n     On entry, the matrix to be inverted.
///                                        On exit, the inverse.
/// \param[out]    work       size n*n     work space for computing the inverse
/// \param[out]    iwork      size 2*n     integer work space
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::matrix_inverse_lu(const int n,
                                                S & det,
                                                S * const A,
                                                S * const work,
                                                int * const iwork) {

  // use some of the integer work space for pivoting during the LU decomposition
  int * const ipiv = iwork;

  // use some of the integer work space for inverting the LU pivoting
  int * const inv_perm = iwork + n;

  // copy the matrix into the workspace
  pcps::xcopy(n*n, A, 1, work, 1);

  // compute the LU decomposition of the matrix
  int info;
  pcps::xgetrf(n, n, work, n, ipiv, info);
  if (info < 0)
    throw pcps::Exception( (boost::format("pcps::xgetrf failed with error code %i in pcps::matrix_inverse_lu") % info).str() );
  if (info > 0)
    throw pcps::SingularMatrixException( (boost::format("U(%i,%i) is exactly zero in pcps::matrix_inverse_lu, thus the matrix is not invertable") % info % info).str() );

  // since FORTRAN counts from 1, subtract 1 from each of the pivot array values
  for (int i = 0; i < n; i++)
    ipiv[i] -= 1;

  // compute the determinant
  det = pcps::unity<S>();
  for (int i = 0; i < n; i++)
    if ( i != ipiv[i] )
      det *= - work[i*n+i];
    else
      det *=   work[i*n+i];
  //if (std::abs(det) < 1.0e-12)
  //  throw pcps::SingularMatrixException( (boost::format("matrix determinant is zero in pcps::matrix_inverse_lu, thus the matrix is not invertable")).str() );

  // determine the inverse of the pivot array's permutation
  for (int i = 0; i < n; i++)
    inv_perm[i] = i;
  for (int i = n-1; i >= 0; i--)
    std::swap(inv_perm[i], inv_perm[ipiv[i]]);

  // compute the inverse of the upper triangle matrix
  for (int i = 0; i < n; i++)
  for (int j = i; j < n; j++) {
    A[i*n+j] = ( i == j ? pcps::unity<S>() : pcps::zero<S>() );
    for (int k = i; k < j; k++)
      A[i*n+j] -= A[i*n+k] * work[j*n+k];  // work indices transposed due to FORTRAN
    A[i*n+j] /= work[j*n+j];
  }

  // compute the inverse of the lower triangle matrix (whose diagonal elements are one and are not stored)
  for (int i = 1; i < n; i++)
  for (int j = i-1; j >= 0; j--) {
    A[i*n+j] = -work[j*n+i];  // work indices transposed due to FORTRAN
    for (int k = j+1; k < i; k++)
      A[i*n+j] -= A[k*n+j] * work[k*n+i];  // work indices transposed due to FORTRAN
  }

  // compute the product of the inverse triangular matrices
  {
    const S * const U_inv = A;
    const S * const L_inv = A;
    for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++) {
      // term with diagonal of L_inv (which equals one)
      work[i*n+j] = ( i <= j ? U_inv[i*n+j] : pcps::zero<S>() );
      // other terms
      for (int k = std::max(i, j+1); k < n; k++)
        work[i*n+j] += U_inv[i*n+k] * L_inv[k*n+j];
    }
  }

  // Multiply by the inverse of the permutation matrix (column reordering)
  // We also transpose the final matrix to account for FORTRAN indexing
  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++)
    A[j*n+i] = work[i*n+inv_perm[j]];

}

// explicitly instantiate the function
template void pcps::matrix_inverse_lu(const int n,
                                      double & det,
                                      double * const A,
                                      double * const work,
                                      int * const iwork);
template void pcps::matrix_inverse_lu(const int n,
                                      std::complex<double> & det,
                                      std::complex<double> * const A,
                                      std::complex<double> * const work,
                                      int * const iwork);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the determinant of a matrix using LU decomposition
///
/// \param[in]       n       The matrix dimension.
/// \param[in,out]   A       On entry, the matrix whose determinant we want.  Overwritten on exit. Size n*n.
/// \param[out]      iwork   Work space for pivoting during the LU decomposition, size n.
///
/// \return the determinant in mantissa/exponent format
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> pcps::MantExp<S> pcps::matrix_determinant_lu(const int n, S * const A, int * const iwork) {

  // compute the LU decomposition of the matrix
  int info;
  pcps::xgetrf(n, n, A, n, iwork, info);
  if (info < 0)
    throw pcps::Exception("pcps::xgetrf failed with error code %i in pcps::matrix_determinant_lu") % info;

  // since FORTRAN counts from 1, subtract 1 from each of the pivot array values
  for (int i = 0; i < n; i++)
    iwork[i] -= 1;

  // compute the determinant
  pcps::MantExp<S> det(pcps::unity<S>());
  for (int i = 0; i < n; i++)
    if ( i != iwork[i] )
      det *= - A[i*n+i];
    else
      det *=   A[i*n+i];

  // return the determinant value
  return det;

}

// explicitly instantiate the function
template pcps::MantExp<double               > pcps::matrix_determinant_lu(const int, double               * const, int * const);
template pcps::MantExp<std::complex<double> > pcps::matrix_determinant_lu(const int, std::complex<double> * const, int * const);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Constructs the VZ and UZT matrices.
///
/// \param[in]       n                           The dimension of A.
/// \param[in]       m                           The number of rows in U and V.
/// \param[in]       Z       size n*n            The inverse of A.             Z_{i,j} =   Z[i*n+j]      i < n, j < n
/// \param[in]       U       size m*n            The matrix U.                 U_{i,j} =   U[i*n+j]      i < m, j < n
/// \param[in]       V       size m*n            The matrix V.                 V_{i,j} =   V[i*n+j]      i < m, j < n
/// \param[in,out]   UZT     size m*n            The product of U and Z^T.   UZT_{i,j} = UZT[i*n+j]      i < m, j < n
/// \param[in,out]   VZ      size m*n            The product of V and Z.      VZ_{i,j} =  VZ[i*n+j]      i < m, j < n
/// \param[in,out]   Umask   size m              For each row of UZT, whether or not that inverse times vector
///                                              product is already known and stored in UZT.
/// \param[in,out]   Vmask   size m              For each row of VZ, whether or not that inverse times vector
///                                              product is already known and stored in VZ.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::matrix_math_build_VZ_UZT(const int n,
                                                       const int m,
                                                       const S * const Z,
                                                       const S * const U,
                                                       const S * const V,
                                                       S * const UZT,
                                                       S * const VZ,
                                                       int * const Umask,
                                                       int * const Vmask) {

  for (int i = 0; i < m; i++) {

    if ( Vmask[i] == 0 ) {
      pcps::xgemm('N', 'T', 1, n, n, pcps::unity<S>(), V + i*n, 1, Z, n, pcps::zero<S>(), VZ + i*n, 1);
      Vmask[i] = 1;
    }

    if ( Umask[i] == 0 ) {
      pcps::xgemm('N', 'N', 1, n, n, pcps::unity<S>(), U + i*n, 1, Z, n, pcps::zero<S>(), UZT + i*n, 1);
      Umask[i] = 1;
    }

  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Constructs the Q matrix.
///
/// \param[in]       n                           The dimension of A.
/// \param[in]       m                           The number of rows in U and V.
/// \param[in]       Z       size n*n            The inverse of A.             Z_{i,j} =   Z[i*n+j]      i < n, j < n
/// \param[in]       U       size m*n            The matrix U.                 U_{i,j} =   U[i*n+j]      i < m, j < n
/// \param[in]       V       size m*n            The matrix V.                 V_{i,j} =   V[i*n+j]      i < m, j < n
/// \param[in,out]   UZT     size m*n            The product of U and Z^T.   UZT_{i,j} = UZT[i*n+j]      i < m, j < n
/// \param[in,out]   VZ      size m*n            The product of V and Z.      VZ_{i,j} =  VZ[i*n+j]      i < m, j < n
/// \param[in,out]   Q       size m*m            Q_{i,j} = Q[i*m+j] = sum_{k,l} V_{i,k} Z_{k,l} U_{j,l}  i < m, j < m
/// \param[in,out]   Umask   size m              For each row of UZT, whether or not that inverse times vector
///                                              product is already known and stored in UZT.
/// \param[in,out]   Vmask   size m              For each row of VZ, whether or not that inverse times vector
///                                              product is already known and stored in VZ.
/// \param[in,out]   Qmask   size m*m            For each element of Q, whether or not that element is already known
///                                              and stored in Q.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::matrix_math_build_Q(const int n,
                                                  const int m,
                                                  const S * const Z,
                                                  const S * const U,
                                                  const S * const V,
                                                  S * const UZT,
                                                  S * const VZ,
                                                  S * const Q,
                                                  int * const Umask,
                                                  int * const Vmask,
                                                  int * const Qmask) {

  for (int i = 0; i < m; i++)
  for (int j = 0; j < m; j++) {

    // if we don't already know the element of Q, compute it
    if ( Qmask[i*m+j] == 0 ) {

      // if we know the row of VZ then use that
      if ( Vmask[i] != 0 )
        Q[i*m+j] = pcps::xdot(n, VZ + i*n, 1, U + j*n, 1);

      // otherwise, if we know the row of UZT then use that
      else if ( Umask[j] != 0 )
        Q[i*m+j] = pcps::xdot(n, V + i*n, 1, UZT + j*n, 1);

      // otherwise, compute the row of VZ and use that
      else {
        pcps::xgemm('N', 'T', 1, n, n, pcps::unity<S>(), V + i*n, 1, Z, n, pcps::zero<S>(), VZ + i*n, 1);
        Vmask[i] = 1;
        Q[i*m+j] = pcps::xdot(n, VZ + i*n, 1, U + j*n, 1);
      }

      // mark that we now know the element of Q
      Qmask[i*m+j] = 1;

    }

  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Computes the determinant ratio det(A + U^T C V) / det(A).
///
/// \param[in]       n                    The dimension of A and B.
/// \param[in]       m                    The number of rows in U and V.
/// \param[in]       Z       size n*n     The inverse of A.             Z_{i,j} =   Z[i*n+j]      i < n, j < n
/// \param[in]       U       size m*n     The matrix U.                 U_{i,j} =   U[i*n+j]      i < m, j < n
/// \param[in]       V       size m*n     The matrix V.                 V_{i,j} =   V[i*n+j]      i < m, j < n
/// \param[in]       C       size m*m     The matrix C.                 C_{i,j} =   C[i*m+j]      i < m, j < m
/// \param[in,out]   UZT     size m*n     The product of U and Z^T.   UZT_{i,j} = UZT[i*n+j]      i < m, j < n
/// \param[in,out]   VZ      size m*n     The product of V and Z.      VZ_{i,j} =  VZ[i*n+j]      i < m, j < n
/// \param[in,out]   Q       size m*m     Q_{i,j} = Q[i*m+j] = sum_{k,l} V_{i,k} Z_{k,l} U_{j,l}  i < m, j < m
/// \param[in,out]   Umask   size m       For each row of UZT, whether or not that inverse times vector
///                                       product is already known and stored in UZT.
/// \param[in,out]   Vmask   size m       For each row of VZ, whether or not that inverse times vector
///                                       product is already known and stored in VZ.
/// \param[in,out]   Qmask   size m*m     For each element of Q, whether or not that element is already known
///                                       and stored in Q.
/// \param[out]      work    size 2*m*m   work space
/// \param[out]      iwork   size 2*m     work space
/// \param[in]       c_eye                true if C is the identity matrix, false otherwise
///
/// \return  the ratio    det(A + U^T C V) / det(A) = det( C^(-1) + V Z U^T )
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> S pcps::matrix_determinant_lemma(const int n,
                                                    const int m,
                                                    const S * const Z,
                                                    const S * const U,
                                                    const S * const V,
                                                    const S * const C,
                                                    S * const UZT,
                                                    S * const VZ,
                                                    S * const Q,
                                                    int * const Umask,
                                                    int * const Vmask,
                                                    int * const Qmask,
                                                    S * const work,
                                                    int * const iwork,
                                                    const bool c_eye) {

  // build Q
  pcps::matrix_math_build_Q(n, m, Z, U, V, UZT, VZ, Q, Umask, Vmask, Qmask);

  // // create the matrix I + Q in the workspace
  // pcps::xscal(m*m, pcps::zero<S>(), work, 1);
  // for (int i = 0; i < m; i++)
  //   work[i*m+i] = pcps::unity<S>();
  // pcps::xaxpy(m*m, pcps::unity<S>(), Q, 1, work, 1);

  S * const x = work + m*m;

  // if C is the identity matrix
  if ( c_eye ) {

    // create the matrix I + Q
    pcps::xcopy(m*m, Q, 1, x, 1);
    x[0*3+0] += pcps::unity<S>();
    x[1*3+1] += pcps::unity<S>();
    x[2*3+2] += pcps::unity<S>();

  // otherwise
  } else {

    // invert C
    pcps::xcopy(m*m, C, 1, work, 1);
    S det_val;
    pcps::matrix_inverse_lu(m, det_val, work, x, iwork);

    // create the matrix C^(-1) + Q
    pcps::xcopy(m*m, work, 1, x, 1);
    pcps::xaxpy(m*m, pcps::unity<S>(), Q, 1, x, 1);

  }

  // get variable to hold return value
  S retval;

  // dirct determinant calculation for 1x1 matrices
  if (m == 1)
    retval = x[0];

  // dirct determinant calculation for 2x2 matrices
  else if (m == 2)
    retval = ( x[0*2+0] * x[1*2+1] - x[1*2+0] * x[0*2+1] );

  // dirct determinant calculation for 3x3 matrices
  else if (m == 3)
    retval = (   x[0*3+0] * ( x[1*3+1] * x[2*3+2] - x[2*3+1] * x[1*3+2] )
               - x[0*3+1] * ( x[1*3+0] * x[2*3+2] - x[2*3+0] * x[1*3+2] )
               + x[0*3+2] * ( x[1*3+0] * x[2*3+1] - x[2*3+0] * x[1*3+1] ) );

  // determinant calculation for general mxm matrices
  else
    retval = pcps::matrix_determinant_lu(m, x, iwork).get();

  // return the determinant of C^(-1) + Q, which is the ratio we want
  return retval;

}

template double pcps::matrix_determinant_lemma(const int n,
                                               const int m,
                                               const double * const Z,
                                               const double * const U,
                                               const double * const V,
                                               const double * const C,
                                               double * const UZT,
                                               double * const VZ,
                                               double * const Q,
                                               int * const Umask,
                                               int * const Vmask,
                                               int * const Qmask,
                                               double * const work,
                                               int * const iwork,
                                               const bool c_eye);

template std::complex<double> pcps::matrix_determinant_lemma(const int n,
                                                             const int m,
                                                             const std::complex<double> * const Z,
                                                             const std::complex<double> * const U,
                                                             const std::complex<double> * const V,
                                                             const std::complex<double> * const C,
                                                             std::complex<double> * const UZT,
                                                             std::complex<double> * const VZ,
                                                             std::complex<double> * const Q,
                                                             int * const Umask,
                                                             int * const Vmask,
                                                             int * const Qmask,
                                                             std::complex<double> * const work,
                                                             int * const iwork,
                                                             const bool c_eye);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Given the inverse Z of a matrix A, updates Z as
///          Z  -->  ( A + U^T C V )^(-1)  =  Z - Z U^T ( C^(-1) + V Z U^T )^(-1) V Z
///
/// \param[in]       n                           The dimension of A.
/// \param[in]       m                           The number of rows in U and V.
/// \param[in,out]   Z       size n*n            On entry, the inverse of A.   Z_{i,j} =  Z[i*n+j]       i < n, j < n
///                                              On exit, the inverse of A + U^T C V.
/// \param[in]       U       size m*n            The matrix U.                 U_{i,j} =   U[i*n+j]      i < m, j < n
/// \param[in]       V       size m*n            The matrix V.                 V_{i,j} =   V[i*n+j]      i < m, j < n
/// \param[in]       C       size m*m            The matrix C.                 C_{i,j} =   C[i*m+j]      i < m, j < m
/// \param[in,out]   UZT     size m*n            The product of U and Z^T.   UZT_{i,j} = UZT[i*n+j]      i < m, j < n
/// \param[in,out]   VZ      size m*n            The product of V and Z.      VZ_{i,j} =  VZ[i*n+j]      i < m, j < n
/// \param[in,out]   Q       size m*m            Q_{i,j} = Q[i*m+j] = sum_{k,l} V_{i,k} Z_{k,l} U_{j,l}  i < m, j < m
/// \param[in,out]   Umask   size m              For each row of UZT, whether or not that inverse times vector
///                                              product is already known and stored in UZT.
/// \param[in,out]   Vmask   size m              For each row of VZ, whether or not that inverse times vector
///                                              product is already known and stored in VZ.
/// \param[in,out]   Qmask   size m*m            For each element of Q, whether or not that element is already known
///                                              and stored in Q.
/// \param[out]      work    size 2*m*m + n*m    work space
/// \param[out]      iwork   size 2*m            work space
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::woodbury_matrix_identity(const int n,
                                                       const int m,
                                                       S * const Z,
                                                       const S * const U,
                                                       const S * const V,
                                                       const S * const C,
                                                       S * const UZT,
                                                       S * const VZ,
                                                       S * const Q,
                                                       int * const Umask,
                                                       int * const Vmask,
                                                       int * const Qmask,
                                                       S * const work,
                                                       int * const iwork) {

  // build VZ, UZT, and Q
  pcps::matrix_math_build_VZ_UZT(n, m, Z, U, V, UZT, VZ, Umask, Vmask);
  pcps::matrix_math_build_Q(n, m, Z, U, V, UZT, VZ, Q, Umask, Vmask, Qmask);

  // invert C
  pcps::xcopy(m*m, C, 1, work, 1);
  S det_val;
  pcps::matrix_inverse_lu(m, det_val, work, work+m*m, iwork);

  // create the matrix C^(-1) + Q
  pcps::xcopy(m*m, work, 1, work + m*m, 1);
  pcps::xaxpy(m*m, pcps::unity<S>(), Q, 1, work + m*m, 1);

  // compute the inverse of C^(-1) + Q
  pcps::matrix_inverse_lu(m, det_val, work + m*m, work, iwork);

//  for (int i = 0; i < n; i++)
//  for (int j = 0; j < n; j++) {
//    S x = pcps::zero<S>();
//    for (int p = 0; p < m; p++)
//    for (int q = 0; q < m; q++)
//      x += UZT[p*n+i] * work[m*m+p*m+q] * VZ[q*n+j];
//    Z[i*n+j] -= x;
//  }

  // compute the product ( C^(-1) + Q )^(-1) V Z
  pcps::xgemm('N', 'N', n, m, m, pcps::unity<S>(), VZ, n, work + m*m, m, pcps::zero<S>(), work + 2*m*m, n);

  // update the inverse of A
  pcps::xgemm('N', 'T', n, n, m, -pcps::unity<S>(), work + 2*m*m, n, UZT, n, pcps::unity<S>(), Z, n);

}

template void pcps::woodbury_matrix_identity(const int n,
                                             const int m,
                                             double * const Z,
                                             const double * const U,
                                             const double * const V,
                                             const double * const C,
                                             double * const UZT,
                                             double * const VZ,
                                             double * const Q,
                                             int * const Umask,
                                             int * const Vmask,
                                             int * const Qmask,
                                             double * const work,
                                             int * const iwork);

template void pcps::woodbury_matrix_identity(const int n,
                                             const int m,
                                             std::complex<double> * const Z,
                                             const std::complex<double> * const U,
                                             const std::complex<double> * const V,
                                             const std::complex<double> * const C,
                                             std::complex<double> * const UZT,
                                             std::complex<double> * const VZ,
                                             std::complex<double> * const Q,
                                             int * const Umask,
                                             int * const Vmask,
                                             int * const Qmask,
                                             std::complex<double> * const work,
                                             int * const iwork);
